from reinforce_brain import PolicyGradient
import gym
import matplotlib.pyplot  as plt



RENDER = False  # 不显示画面
DISPLAY_REWARD_THRESHOLD = 400  # 当回合总reward大于400时显示模拟窗口

env = gym.make('CartPole-v0')
env = env.unwrapped  # 取消限制
env.seed(1)  # 普通的Policy gradient方法,似的回合的variance比较大,所以选择逸一个好一点的随机种子

print(env.action_space)  # 显示可用的action
print(env.observation_space)  # 显示可用state的observation
print(env.observation_space.high)  # 显示observation最高值
print(env.observation_space.low)  # 显示observation最低值
# observation = env.reset()
# print(observation)

# 定义
RL = PolicyGradient(
    n_actions=env.action_space.n,
    n_features=env.observation_space.shape[0],
    learning_rate=0.02,
    reward_decay=0.99,
)

# agent执行完一整幕游戏后进行更新
for i_episode in range(3000):
    observation = env.reset()
    while True:
        if RENDER: env.render()
        # 选择动作
        action = RL.choose_action(observation)
        # 执行动作获得观测值
        observation_, reward, done, _ = env.step(action)
        # 存储这一回合的transition
        RL.store_transition(observation, action, reward)

        # 在一幕结束后进行更新
        if done:
            ep_rs_sum = sum(RL.ep_rs)

            if 'running_reward' not in globals():
                running_reward = ep_rs_sum
            else:
                running_reward = running_reward * 0.99 + ep_rs_sum * 0.01
            if running_reward > DISPLAY_REWARD_THRESHOLD: RENDER = True     # rendering
            print("episode:", i_episode, "  reward:", int(running_reward))

            vt = RL.learn()
            if i_episode == 0:
                plt.plot(vt)    # plot the episode vt
                plt.xlabel('episode steps')
                plt.ylabel('normalized state-action value')
                plt.show()
            break
        observation = observation_
